import os
import numpy as np
from matplotlib import pyplot as plt

import config
import getGLP
import getBasis
import Error

# -------------------------------------------
# Utilities: reconstruct & plot at a given time
# -------------------------------------------
def plot_solution(uh, t, weight, lambda_, Xc, hx1, title_prefix="DG solution", ax=None, filename=None):
    """
    Plot numerical solution at the current time using Gauss points in each cell.

    Parameters
    ----------
    uh : ndarray, shape (Nx, dimPk)
        Coefficients per cell.
    t : float
        Current time.
    weight, lambda_ : arrays
        Gauss weights and points on reference cell [-1,1].
    Xc : ndarray, shape (Nx,)
        Cell centers.
    hx1 : float
        Half of cell size (mapping x = hx1 * λ + Xc).
    title_prefix : str
        Title prefix for the figure.
    ax : matplotlib.axes.Axes or None
        If provided, draw on this axes; otherwise create a new figure.
    filename : str or None
        If provided, save the figure to this path.
    """
    Nx, dimPk = uh.shape
    NumGLP = len(lambda_)

    # Reconstruct uh at Gauss points of each cell
    uhG = np.zeros((Nx, NumGLP), dtype=float)
    for i in range(Nx):
        # uhG[i, :] = sum_d uh[i, d] * phiG[:, d]
        # (phiG is NumGLP x dimPk on reference cell)
        uhG[i, :] = (getBasis.phiG @ uh[i, :].reshape(-1, 1)).ravel()

    # Physical coordinates for each (cell, gauss) pair
    x_all = np.zeros((Nx, NumGLP), dtype=float)
    for i in range(Nx):
        x_all[i, :] = hx1 * lambda_ + Xc[i]

    # Flatten in cell order to draw one continuous line
    x_plot = x_all.ravel()
    u_plot = uhG.ravel()

    # Sort by x to avoid zig-zagging (optional but nicer)
    sorter = np.argsort(x_plot)
    x_plot = x_plot[sorter]
    u_plot = u_plot[sorter]

    created_fig = False
    if ax is None:
        fig, ax = plt.subplots(figsize=(7, 4))
        created_fig = True

    ax.plot(x_plot, u_plot, linewidth=1.8)
    ax.set_xlabel("x")
    ax.set_ylabel("u_h(x, t)")
    ax.set_title(f"{title_prefix}  (t = {t:.6f})")
    ax.grid(True, linestyle="--", alpha=0.4)

    if filename is not None:
        # Ensure directory exists
        os.makedirs(os.path.dirname(filename), exist_ok=True)
        plt.savefig(filename, dpi=150, bbox_inches="tight")

    if created_fig:
        plt.show()
        plt.close()


# -------------------------------------------
# Spatial operator (same as your main, kept here for encapsulation)
# -------------------------------------------
def Lh(uh: np.ndarray, f=lambda u: u):
    """
    Semi-discrete spatial operator for 1D DG on uniform mesh with Lax-Friedrichs flux.

    Input:
        uh: (Nx, dimPk) cellwise coefficients
        f : flux function supporting ndarray operations
    Output:
        du: (Nx, dimPk)
    """
    Nx = config.Nx
    dimPk = config.dimPk
    NumGLP = config.NumGLP
    weight, _lambda = getGLP.load_GLP(NumGLP)

    hx = config.hx
    bcL = config.bcL
    bcR = config.bcR

    du = np.zeros((Nx, dimPk), dtype=float)
    uhG = np.zeros((Nx, NumGLP), dtype=float)
    uhb = np.vstack([np.zeros((1, dimPk), dtype=float), uh, np.zeros((1, dimPk), dtype=float)])
    uhR = np.zeros((Nx+1, 1), dtype=float)
    uhL = np.zeros((Nx+1, 1), dtype=float)
    f_hat = np.zeros((Nx+1, 1), dtype=float)

    # Periodic boundaries if requested
    if bcL == 1:
        uhb[0, :] = uh[-1, :]
    if bcR == 1:
        uhb[-1, :] = uh[0, :]

    # Step 1: cell interior contributions (at Gauss points)
    for i in range(Nx):
        for d in range(dimPk):
            uhG[i] += uh[i, d] * getBasis.phiG[:, d].T

    for i in range(Nx):
        for d in range(1, dimPk):
            for i1 in range(NumGLP):
                du[i, d] += 0.5 * weight[i1] * f(uhG[i, i1]) * getBasis.phixG[i1, d]

    # Step 2: numerical flux at interfaces
    for i in range(Nx+1):
        for d in range(dimPk):
            uhR[i] += uhb[i, d] * getBasis.phiGR[d]
            uhL[i] += uhb[i+1, d] * getBasis.phiGL[d]

    for i in range(Nx+1):
        uR = uhL[i]
        uL = uhR[i]
        alpha = 1.0
        f_hat[i] = 0.5 * (f(uR) + f(uL) - alpha * (uR - uL))

    # Step 3: surface contributions distributed to each cell
    for i in range(Nx):
        for d in range(dimPk):
            du[i, d] += -1.0 / hx * (getBasis.phiGR[d] * f_hat[i+1] - getBasis.phiGL[d] * f_hat[i])

    # Step 4: divide by diagonal mass matrix
    for d in range(dimPk):
        du[:, d] /= getBasis.mm[d]

    return du


# -------------------------------------------
# Main driver: simulate and plot every step (or every N steps)
# -------------------------------------------
def run_dg_and_plot(record_every=1, save_dir=None, title_prefix="DG solution", flux=lambda u: u):
    """
    Run the DG-RK3 solver and plot the numerical solution over time.

    Parameters
    ----------
    record_every : int
        Plot every `record_every` time steps. Use 1 to plot every step.
    save_dir : str or None
        If provided, figures are saved into this directory with zero-padded indices.
        If None, figures are only shown (not saved).
    title_prefix : str
        Prefix for figure titles.
    flux : callable
        Flux function f(u). Default is linear advection f(u)=u.

    Returns
    -------
    uh : ndarray (Nx, dimPk)
        Final DG coefficients at t = t_end.
    """
    # --- Unpack config ---
    Nx = config.Nx
    k = config.k
    NumGLP = config.NumGLP
    dimPk = config.dimPk
    CFL = config.CFL
    xa = config.xa
    xb = config.xb
    bcL = config.bcL
    bcR = config.bcR
    t_end = config.t_end

    weight, lambda_ = getGLP.load_GLP(NumGLP)
    t = 0.0
    hx = config.hx
    hx1 = config.hx1
    Xc = config.Xc.copy()
    dt = CFL * hx

    # --- L2 projection of initial condition onto DG space ---
    u_real = np.zeros((Nx, NumGLP), dtype=float)
    for i in range(Nx):
        for j in range(NumGLP):
            x = hx1 * lambda_[j] + Xc[i]
            u_real[i, j] = config.u0(x)

    uh = np.zeros((Nx, dimPk), dtype=float)
    for i in range(Nx):
        for d in range(dimPk):
            for i1 in range(NumGLP):
                uh[i, d] += 0.5 * weight[i1] * u_real[i, i1] * getBasis.phiG[i1, d]
    for d in range(dimPk):
        uh[:, d] /= getBasis.mm[d]

    # Initial plot (t=0)
    step = 0
    if record_every > 0 and step % record_every == 0:
        fname = None
        if save_dir is not None:
            fname = os.path.join(save_dir, f"frame_{step:06d}.png")
        plot_solution(uh, t, weight, lambda_, Xc, hx1, title_prefix=title_prefix, filename=fname)

    # --- RK3 time stepping ---
    while t < t_end:
        if t + dt >= t_end:
            dt = t_end - t
            t = t_end
        else:
            t += dt

        # stage I
        du = Lh(uh, f=flux)
        uh1 = uh + dt * du

        # stage II
        du = Lh(uh1, f=flux)
        uh2 = 3.0 / 4.0 * uh + 1.0 / 4.0 * (uh1 + dt * du)

        # stage III
        du = Lh(uh2, f=flux)
        uh = 1.0 / 3.0 * uh + 2.0 / 3.0 * (uh2 + dt * du)

        # plot this step
        step += 1
        if record_every > 0 and step % record_every == 0:
            fname = None
            if save_dir is not None:
                fname = os.path.join(save_dir, f"frame_{step:06d}.png")
            plot_solution(uh, t, weight, lambda_, Xc, hx1, title_prefix=title_prefix, filename=fname)

    # Optionally report L2 error at final time if you have a reference
    # (Here we reuse the initial u_real only for demonstration; replace with exact at t_end if available)
    L2_err = Error.calculate_L2_Error(
        uh=uh,
        u_real=u_real,           # replace with exact u(x, t_end) if you have it
        Nx=Nx,
        NumGLP=NumGLP,
        dimPk=dimPk,
        phiG=getBasis.phiG,
        weight=weight,
        hx1=hx1
    )
    print(f"Final L2 error (vs. initial field on GL points) = {L2_err}")

    return uh